import numpy as np
from functionSim_config import *
from easySample import easySample
import shelve
import torch
import sys
sys.path.append(r"/home/cyw/projects/function_sim_project/basic_script")


class sample_function_embedding():
    """
        获得样本经过模型训练后的各个函数的嵌入值
        input:模型，样本a，样本b
        output:两个样本的图节点的嵌入值
    """

    def __init__(self, model, namea, nameb) -> None:
        """
            namea为样本的MD5值
        """
        # shelve 名字区分大小写,大写的时候生成了一个新的shelve文件
        namea = namea.lower()
        nameb = nameb.lower()

        self.model = torch.load(model)
        sample = easySample()

        sample_a = sample.get_sample(namea, "functionSim")
        sample_b = sample.get_sample(nameb, "functionSim")

        x_adj = torch.tensor([sample_a["adj"]], dtype=torch.float64).to(device)
        x_att = torch.tensor([sample_a["att"]], dtype=torch.float64).to(device)

        x_type = torch.tensor([sample_a["vtype"]],
                              dtype=torch.float64).to(device)
        y_adj = torch.tensor([sample_b["adj"]], dtype=torch.float64).to(device)
        y_att = torch.tensor([sample_b["att"]], dtype=torch.float64).to(device)
        y_type = torch.tensor([sample_b["vtype"]],
                              dtype=torch.float64).to(device)
        score1 = self.model(x_adj, x_att, x_type, y_adj, y_att, y_type)
        print(score1)
        self.get_embedding(x_type[0].tolist(), y_type[0].tolist())

    def softmax(self, scores):
        exp_scores = np.exp(scores)
        probabilities = exp_scores / np.sum(exp_scores, axis=0)
        return probabilities

    def get_embedding(self, x_type, y_type):
        """
            获得保存的样本的函数嵌入
            并计算两两间的相似值
            输入为样本的类型，用于找出不同的函数
        """
        functionPath = "/home/cyw/projects/function_sim_project/all_data/functionEmbedding"
        with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/functionEmbedding/{}".format("sample_A")) as file:
            sampleA = file["embedding"][0]
            file.close()
        with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/functionEmbedding/{}".format("sample_B")) as file:
            sampleB = file["embedding"][0]
            file.close()
        m, n = len(sampleA), len(sampleB)
        print("样本A： {}个节点，样本B： {}个节点".format(m, n))
        # 节点太多了37个*63个，展示不了一点    和之前构思的情况不太一样
        # 看一下ida中的情况，只展示内部函数的部分怎么样？第一个样本6个节点，第二个样本15个节点，一共90条线，勉强可以画出来
        # 如何解释----》外部函数能通过函数名精准匹配，我们这里只展示不能准确匹配的内部函数之间的对应关系
        local_a, local_b = 0, 0
        # ！！！这里默认本地函数在别的函数前面，这两个样例符合，别的样例要注意是否正确！！！
        for i in range(len(x_type)):
            if int(x_type[i][0]) == 1:
                local_a += 1
        for i in range(len(y_type)):
            if int(y_type[i][0]) == 1:
                local_b += 1
        # 计算两两间的相似值
        originRes = [[0]*local_b for i in range(local_a)]
        for i in range(local_a):
            tempSum = 0
            for j in range(local_b):
                cosine_similarity = round(
                    float(torch.nn.functional.cosine_similarity(sampleA[i], sampleB[j], dim=0)), 2)
                originRes[i][j] = cosine_similarity
                tempSum += cosine_similarity
            for j in range(local_b):
                temp = round(originRes[i][j]/tempSum, 3)
                originRes[i][j] = temp
        print("相似矩阵：")
        for i in range(local_a):
            for j in range(local_b):
                print("%7s" % (str(originRes[i][j])), end="")
            print("")


if __name__ == "__main__":
    # !!!!!!!!切换的时候配置信息中的设置也需要改变！！！！！！！！！！  难得写代码了，用的时候手动注释一下！！！！！！
    functionSimModel = "/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge_best_4_8.pth"

    sample_a = "A7B483A4E74853C0836DE21EAFD5DD83"
    sample_b = "74AA01EB4412C873E732376168F3B772"

    a = sample_function_embedding(functionSimModel, sample_a, sample_b)
